import { Model, ModelMetadataExtensions } from '../../Core/Model';
// import { Tensor } from '../../Core/Tensor';
import { TensorNames } from './TensorNames';
import { TensorProxy, TensorType } from './TensorProxy';

export interface FailedCheck {
    message: string;
}

export class ModelExtensions {
    static getInputNames(model: Model): TensorProxy[] {
        const names: TensorProxy[] = [];
        if (model == null) return names;
        for (const input of model.inputs) {
            const tensorProxy = new TensorProxy();
            tensorProxy.name = input.name;
            tensorProxy.valueType = TensorType.FloatingPoint;
            tensorProxy.shape = input.shape.map((i) => i as number);
            names.push(tensorProxy);
        }
        names.sort((a, b) => a.name.localeCompare(b.name));
        return names;
    }

    static getVersion(model: Model): number {
        return ModelMetadataExtensions.getTensorByName(model, TensorNames.versionNumber).getCache(0);
    }

    static getInputTensors(model: Model): TensorProxy[] {
        const tensors: TensorProxy[] = [];
        if (model == null) return tensors;
        for (const input of model.inputs) {
            const tensor = new TensorProxy();
            tensor.name = input.name;
            tensor.valueType = TensorType.FloatingPoint;
            tensor.shape = input.shape.map((i) => i as number);
            tensors.push(tensor);
        }
        tensors.sort((a, b) => a.name.localeCompare(b.name));
        return tensors;
    }

    static getNumVisualInputs(model: Model): number {
        let count = 0;
        if (model == null) return count;
        for (const input of model.inputs) {
            if (input.name.startsWith(TensorNames.visualObservationPlaceholderPrefix)) {
                count++;
            }
        }
        return count;
    }

    static getOutputNames(model: Model, deterministicInference = false): string[] {
        const names: string[] = [];
        if (model == null) return names;
        if (this.hasContinuousOutputs(model, deterministicInference)) {
            names.push(this.continuousOutputName(model, deterministicInference));
        }
        if (this.hasDiscreteOutputs(model, deterministicInference)) {
            names.push(this.discreteOutputName(model, deterministicInference));
        }
        const memory = ModelMetadataExtensions.getTensorByName(model, TensorNames.memorySize);
        if (memory) {
            names.push(TensorNames.recurrentOutput);
        }
        names.sort((a, b) => a.localeCompare(b));
        return names;
    }

    static hasContinuousOutputs(model: Model, deterministicInference = false): boolean {
        if (model == null) return false;
        if (!this.supportsContinuousAndDiscrete(model)) {
            return (
                ModelMetadataExtensions.getTensorByName(model, TensorNames.isContinuousControlDeprecated).getCache(0) >
                0
            );
        } else {
            const hasStochasticOutput =
                !deterministicInference && model.outputs.includes(TensorNames.continuousActionOutput);
            const hasDeterministicOutput =
                deterministicInference && model.outputs.includes(TensorNames.deterministicContinuousActionOutput);
            return (
                (hasStochasticOutput || hasDeterministicOutput) &&
                ModelMetadataExtensions.getTensorByName(model, TensorNames.continuousActionOutputShape).getCache(0) > 0
            );
        }
    }

    static continuousOutputSize(model: Model): number {
        if (model == null) return 0;
        if (!ModelExtensions.supportsContinuousAndDiscrete(model)) {
            return ModelMetadataExtensions.getTensorByName(model, TensorNames.isContinuousControlDeprecated)?.getCache(
                0
            ) > 0
                ? ModelMetadataExtensions.getTensorByName(model, TensorNames.actionOutputShapeDeprecated).getCache(0)
                : 0;
        } else {
            const continuousOutputShape = ModelMetadataExtensions.getTensorByName(
                model,
                TensorNames.continuousActionOutputShape
            );
            return continuousOutputShape == null ? 0 : continuousOutputShape.getCache(0);
        }
    }

    static continuousOutputName(model: Model, deterministicInference = false): string {
        if (model == null) return '';
        if (!this.supportsContinuousAndDiscrete(model)) {
            return TensorNames.actionOutputDeprecated;
        } else {
            return deterministicInference
                ? TensorNames.deterministicContinuousActionOutput
                : TensorNames.continuousActionOutput;
        }
    }

    static hasDiscreteOutputs(model: Model, deterministicInference = false): boolean {
        if (model == null) return false;
        if (!this.supportsContinuousAndDiscrete(model)) {
            return (
                ModelMetadataExtensions.getTensorByName(model, TensorNames.isContinuousControlDeprecated).getCache(0) ==
                0
            );
        } else {
            const hasStochasticOutput =
                !deterministicInference && model.outputs.includes(TensorNames.discreteActionOutput);
            const hasDeterministicOutput =
                deterministicInference && model.outputs.includes(TensorNames.deterministicDiscreteActionOutput);
            return (hasStochasticOutput || hasDeterministicOutput) && this.discreteOutputSize(model) > 0;
        }
    }

    static discreteOutputSize(model: Model): number {
        if (model == null) return 0;
        if (!this.supportsContinuousAndDiscrete(model)) {
            return ModelMetadataExtensions.getTensorByName(model, TensorNames.isContinuousControlDeprecated).getCache(
                0
            ) > 0
                ? 0
                : ModelMetadataExtensions.getTensorByName(model, TensorNames.actionOutputShapeDeprecated).getCache(0);
        } else {
            const discreteOutputShape = ModelMetadataExtensions.getTensorByName(
                model,
                TensorNames.discreteActionOutputShape
            );
            if (discreteOutputShape == null) {
                return 0;
            } else {
                let result = 0;
                for (let i = 0; i < discreteOutputShape.length; i++) {
                    result += discreteOutputShape.getCache(i);
                }
                return result;
            }
        }
    }

    static discreteOutputName(model: Model, deterministicInference = false): string {
        if (model == null) return '';
        if (!this.supportsContinuousAndDiscrete(model)) {
            return TensorNames.actionOutputDeprecated;
        } else {
            return deterministicInference
                ? TensorNames.deterministicDiscreteActionOutput
                : TensorNames.discreteActionOutput;
        }
    }

    static supportsContinuousAndDiscrete(model: Model): boolean {
        return (
            model == null ||
            model.outputs.includes(TensorNames.continuousActionOutput) ||
            model.outputs.includes(TensorNames.discreteActionOutput)
        );
    }

    // static checkExpectedTensors(
    //     model: Model,
    //     failedModelChecks: FailedCheck[],
    //     deterministicInference = false
    // ): boolean {
    //     const modelApiVersionTensor = model.GetTensorByName(TensorNames.VersionNumber);
    //     if (modelApiVersionTensor == null) {
    //         failedModelChecks.push({
    //             message: `Required constant "${TensorNames.VersionNumber}" was not found in the model file.`,
    //         });
    //         return false;
    //     }

    //     const memorySizeTensor = model.GetTensorByName(TensorNames.MemorySize);
    //     if (memorySizeTensor == null) {
    //         failedModelChecks.push({
    //             message: `Required constant "${TensorNames.MemorySize}" was not found in the model file.`,
    //         });
    //         return false;
    //     }

    //     if (
    //         !model.outputs.includes(TensorNames.ActionOutputDeprecated) &&
    //         !model.outputs.includes(TensorNames.ContinuousActionOutput) &&
    //         !model.outputs.includes(TensorNames.DiscreteActionOutput) &&
    //         !model.outputs.includes(TensorNames.DeterministicContinuousActionOutput) &&
    //         !model.outputs.includes(TensorNames.DeterministicDiscreteActionOutput)
    //     ) {
    //         failedModelChecks.push({
    //             message: 'The model does not contain any Action Output Node.',
    //         });
    //         return false;
    //     }

    //     if (!model.SupportsContinuousAndDiscrete()) {
    //         if (model.GetTensorByName(TensorNames.ActionOutputShapeDeprecated) == null) {
    //             failedModelChecks.push({
    //                 message: 'The model does not contain any Action Output Shape Node.',
    //             });
    //             return false;
    //         }
    //         if (model.GetTensorByName(TensorNames.IsContinuousControlDeprecated) == null) {
    //             failedModelChecks.push({
    //                 message: `Required constant "${TensorNames.IsContinuousControlDeprecated}" was not found in the model file. This is only required for model that uses a deprecated model format.`,
    //             });
    //             return false;
    //         }
    //     } else {
    //         if (model.outputs.includes(TensorNames.ContinuousActionOutput)) {
    //             if (model.GetTensorByName(TensorNames.ContinuousActionOutputShape) == null) {
    //                 failedModelChecks.push({
    //                     message:
    //                         'The model uses continuous action but does not contain Continuous Action Output Shape Node.',
    //                 });
    //                 return false;
    //             } else if (!model.hasContinuousOutputs(deterministicInference)) {
    //                 const actionType = deterministicInference ? 'deterministic' : 'stochastic';
    //                 const actionName = deterministicInference ? 'Deterministic' : '';
    //                 failedModelChecks.push({
    //                     message: `The model uses ${actionType} inference but does not contain ${actionName} Continuous Action Output Tensor. Uncheck 'Deterministic inference' flag.`,
    //                 });
    //                 return false;
    //             }
    //         }

    //         if (model.outputs.includes(TensorNames.DiscreteActionOutput)) {
    //             if (model.GetTensorByName(TensorNames.DiscreteActionOutputShape) == null) {
    //                 failedModelChecks.push({
    //                     message:
    //                         'The model uses discrete action but does not contain Discrete Action Output Shape Node.',
    //                 });
    //                 return false;
    //             } else if (!model.hasDiscreteOutputs(deterministicInference)) {
    //                 const actionType = deterministicInference ? 'deterministic' : 'stochastic';
    //                 const actionName = deterministicInference ? 'Deterministic' : '';
    //                 failedModelChecks.push({
    //                     message: `The model uses ${actionType} inference but does not contain ${actionName} Discrete Action Output Tensor. Uncheck 'Deterministic inference' flag.`,
    //                 });
    //                 return false;
    //             }
    //         }
    //     }
    //     return true;
    // }
}
